import streamlit as st
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import numpy as np
import cv2
from matplotlib import rcParams
import io
from PIL import Image
import warnings
warnings.filterwarnings("ignore")
# 设置中文字体
rcParams['font.sans-serif'] = ['SimHei']
rcParams['axes.unicode_minus'] = False


def read_image(uploaded_file):
    image = Image.open(uploaded_file)
    image = np.array(image)
    # 确保图像是RGB格式
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    elif image.shape[2] == 4:
        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
    return image


def color_threshold(image, color='white', thresh=200):
    """颜色阈值筛选"""
    color_select = np.copy(image)

    # 预定义的颜色阈值
    color_thresholds = {
        'white': {'red': thresh, 'green': thresh, 'blue': thresh},
        'yellow': {'red': thresh, 'green': thresh, 'blue': 0},
        'gray': {'red': thresh // 2, 'green': thresh // 2, 'blue': thresh // 2}
    }

    selected_thresh = color_thresholds.get(color, color_thresholds['white'])

    thresholds = (image[:, :, 0] < selected_thresh['red']) \
                 | (image[:, :, 1] < selected_thresh['green']) \
                 | (image[:, :, 2] < selected_thresh['blue'])

    color_select[thresholds] = [0, 0, 0]
    return color_select


def canny_edge_detection(image, kernel_size=5, low_thresh=50, high_thresh=150):
    """Canny 边缘检测"""
    gray_img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    gray_img = cv2.cvtColor(gray_img, cv2.COLOR_BGR2GRAY)
    blur_gray = cv2.GaussianBlur(gray_img, (kernel_size, kernel_size), 0)
    edges = cv2.Canny(blur_gray, low_thresh, high_thresh)
    return edges


def hough_transform(edges, original_image, line_color='red', rho=1, theta=np.pi / 180,
                    threshold=15, min_line_length=40, max_line_gap=20):
    """霍夫变换检测直线"""
    lines = cv2.HoughLinesP(edges, rho, theta, threshold, np.array([]),
                            min_line_length, max_line_gap)

    line_image = np.copy(original_image)

    # 预定义的线条颜色
    color_dict = {
        'red': (255, 0, 0),
        'yellow': (255, 255, 0),
        'green': (0, 255, 0),
        'blue': (0, 0, 255)
    }

    selected_color = color_dict.get(line_color, color_dict['red'])

    if lines is not None:
        for line in lines:
            for x1, y1, x2, y2 in line:
                cv2.line(line_image, (x1, y1), (x2, y2), selected_color, 5)

    return line_image


def overlay_lines(original_image, line_image, alpha=0.8, beta=1, gamma=0):
    """将检测到的线条叠加到原始图像上"""
    return cv2.addWeighted(original_image, alpha, line_image, beta, gamma)


def main():
    st.set_page_config(page_title="车道线检测系统", layout="wide")

    # 页面标题
    st.title("🚗 车道线检测系统")

    # 侧边栏参数设置
    st.sidebar.header("参数设置")

    # 颜色选择
    threshold_color = st.sidebar.selectbox(
        "选择要检测的颜色",
        ['white', 'yellow', 'gray'],
        index=0
    )

    line_color = st.sidebar.selectbox(
        "选择标记线条的颜色",
        ['red', 'yellow', 'green', 'blue'],
        index=0
    )

    # 阈值调整
    thresh = st.sidebar.slider("颜色阈值", 100, 255, 200)

    # 文件上传
    uploaded_file = st.file_uploader("上传待检测的图片", type=['jpg', 'jpeg', 'png'])

    if uploaded_file is not None:
        # 创建三列布局
        col1, col2, col3 = st.columns(3)

        # 读取并显示原始图像
        image = read_image(uploaded_file)
        col1.header("原始图像")
        col1.image(image, use_container_width=True)

        # 颜色阈值筛选
        color_filtered = color_threshold(image, threshold_color, thresh)
        col2.header("颜色筛选结果")
        col2.image(color_filtered, use_container_width=True)

        # 边缘检测
        edges = canny_edge_detection(color_filtered)
        col3.header("边缘检测结果")
        col3.image(edges, use_container_width=True)

        # 霍夫变换和结果显示
        st.header("最终检测结果")
        line_image = hough_transform(edges, image, line_color)
        result_image = overlay_lines(image, line_image)
        st.image(result_image, use_container_width=True)


if __name__ == "__main__":
    main()